import logging
from contextlib import contextmanager

from langgraph.checkpoint.postgres import PostgresSaver
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from app.config.settings import Setting

logger = logging.getLogger(__name__)


class PgManager:
    def __init__(self):
        self.__db_url = f"postgresql://{Setting.pg_user}:{Setting.pg_password}@{Setting.pg_host}:{Setting.pg_port}/{Setting.pg_database}?sslmode=disable"
        # 注意这里的链接字符串稍微有些不同，避免sqlalchemy默认会使用psycopg2版本而出现错误
        self.engine = create_engine(
            f"postgresql+psycopg://{Setting.pg_user}:{Setting.pg_password}@{Setting.pg_host}:{Setting.pg_port}/{Setting.pg_database}?sslmode=disable")
        self.session_func = sessionmaker(bind=self.engine, autoflush=False, autocommit=False)
        logger.warning(f"Connecting to database with URL: {self.__db_url}")

    @contextmanager
    def get_db(self):
        db = self.session_func()
        try:
            yield db
        finally:
            db.close()

    @staticmethod
    def get_db_static():
        """静态方法版本，用于FastAPI依赖注入"""
        from app.config.database_config import PgManager
        manager = PgManager()
        with manager.get_db() as db:
            yield db

    def create_saver(self):
        """
        创建一个PostgresSaver实例，用于保存和恢复状态。
        """
        return PostgresSaver.from_conn_string(self.__db_url)
